/*
Copyright 2020 - 2023 Esri

Licensed under the Apache License, Version 2.0 (the "License"); you may not use
this file except in compliance with the License. You may obtain a copy of
the License at http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software distributed
under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
CONDITIONS OF ANY KIND, either express or implied. See the License for the
specific language governing permissions and limitations under the License.

For additional information, contact:
Environmental Systems Research Institute, Inc.
Attn: Contracts Dept
380 New York Street
Redlands, California, USA 92373
email: contracts@esri.com
*/

#include "i3s/i3s_writer.h"
#include "utils/utl_png.h"
#include "utils/utl_geom.h"
#include "utils/utl_i3s_resource_defines.h"
#include <iostream>
#include <vector>
#include <algorithm>
#include <cmath>
#include <cstdlib>
#include <string>
#include <osg/Material>
#include <osg/PagedLOD>
#include <osg/Math>
#include <osg/Node>
#include <osg/Texture2D>
#include <osgDB/WriteFile>
#include <osgDB/ReadFile>
#include <filesystem>
namespace stdfs = std::filesystem;

using i3slib::utl::Vec2i;
using i3slib::utl::Vec2f;
using i3slib::utl::Vec2d;
using i3slib::utl::Vec3f;
using i3slib::utl::Vec3d;

namespace
{

  constexpr double c_pi = 3.14159265358979323846; // M_PI;
  constexpr double c_wgs84_equatorial_radius = 6378137.0; //semi-major 
  constexpr double c_wgs84_polar_radius = 6356752.314245; //semi-minor
  constexpr double c_degrees_per_meter = 180.0 / (c_wgs84_polar_radius * c_pi);

  struct osg_tree {
    double geometricError;
    std::string file_name;
    std::vector<osg_tree> sub_nodes;
  };

  class InfoVisitor : public osg::NodeVisitor
  {
      std::string path;
  public:
      InfoVisitor(std::string _path)
      :osg::NodeVisitor(TRAVERSE_ALL_CHILDREN)
      ,path(_path)
      {}
  
      ~InfoVisitor() {
      }
  
      void apply(osg::Geometry& geometry){
          geometry_array.push_back(&geometry);
          if (auto ss = geometry.getStateSet() ) {
              osg::Texture* tex = dynamic_cast<osg::Texture*>(ss->getTextureAttribute(0, osg::StateAttribute::TEXTURE));
              if (tex) {
                  texture_array.insert(tex);
                  texture_map[&geometry] = tex;
              }
          }
      }
      
      void apply(osg::PagedLOD& node) {
          //std::string path = node.getDatabasePath();
          int n = node.getNumFileNames();
          
          for (size_t i = 1; i < n; i++)
          {
              std::string file_name = path + "/" + node.getFileName(i);
              sub_node_names.push_back(file_name);
          }
          traverse(node);
      }
  
  public:
      std::vector<osg::Geometry*> geometry_array;
      std::set<osg::Texture*> texture_array;
      // 记录 mesh 和 texture 的关系，暂时认为一个模型最多只有一个 texture
      std::map<osg::Geometry*, osg::Texture*> texture_map; 
      std::vector<std::string> sub_node_names;
  };
  struct CS_transformation
  {
    DECL_PTR(CS_transformation);
    virtual bool transform(Vec3d* points, size_t count) = 0;
  };

  class ENU_to_WGS_transformation : public CS_transformation
  {
  public:
    DECL_PTR(ENU_to_WGS_transformation);

    explicit ENU_to_WGS_transformation(const Vec2d& origin) :
      origin_(origin),
      lon_degrees_per_meter_(180.0 / (c_wgs84_equatorial_radius * c_pi * std::cos(origin.y * c_pi / 180.0)))
    {}

    virtual bool transform(Vec3d* points, size_t count) override
    {
      std::transform(points, points + count, points, [this](const Vec3d& p)
      {
        return Vec3d(origin_.x + p.x * lon_degrees_per_meter_, origin_.y + p.y * c_degrees_per_meter, p.z);
      });
      return true;
    }

  private:
    const Vec2d origin_;
    const double lon_degrees_per_meter_;
  };

  std::string get_parent(std::string str) {
    auto p0 = str.find_last_of("/\\");
    if (p0 != std::string::npos)
        return str.substr(0, p0);
    else
        return "";
  }

  double screen_size_to_area(double pixels)
  {
    constexpr double c_pi_over_4 = c_pi * 0.25;
    return pixels * pixels * c_pi_over_4;
  }

  i3slib::i3s::Layer_writer::Var create_writer(const stdfs::path& slpk_path)
  {
    i3slib::i3s::Ctx_properties ctx_props(i3slib::i3s::Max_major_versions({}));
    //i3slib::i3s::set_geom_compression(ctx_props.geom_encoding_support, i3slib::i3s::Geometry_compression::Draco, true);
    //i3slib::i3s::set_gpu_compression(ctx_props.gpu_tex_encoding_support, i3slib::i3s::GPU_texture_compression::ETC_2, true);
    auto writer_context = i3slib::i3s::create_i3s_writer_context(ctx_props);

    i3slib::i3s::Layer_meta meta;
    meta.type = i3slib::i3s::Layer_type::Mesh_IM;
    meta.name = i3slib::utl::to_string(slpk_path.stem().generic_u8string());
    meta.desc = "Generated with osgb2slpk";
    meta.sr.wkid = 4326;
    meta.uid = meta.name;
    meta.normal_reference_frame = i3slib::i3s::Normal_reference_frame::Not_set;

    // Temp hack to make the output SLPKs be exactly the same in different runs on the same input.
    // TODO: add command line parameter for the timestamp?
    meta.timestamp = 1;

    std::unique_ptr<i3slib::i3s::Layer_writer> writer(
      i3slib::i3s::create_mesh_layer_builder(writer_context, slpk_path));

    if (writer)
      writer->set_layer_meta(meta);
    return writer;
  }

  Vec2f clamp_uv(const Vec2f& uv)
  {
    return
    {
      std::clamp(uv.x, 0.01f, 0.99f),
      std::clamp(uv.y, 0.01f, 0.99f)
    };
  }
  // Extracts a square fragment of given size at a given offset from a square texture.
  // Inputs: src_data is raw RGB, src_size, start and size are in pixels.
  // The output texture is PNG-compressed.
  bool extract_texture_fragment(
    osg::Texture2D* tex2D,
    i3slib::i3s::Texture_buffer& texture_buffer)
  {
    osg::Image* image = tex2D->getImage();
    int width = image->s();
    int height = image->t();

    std::vector<uint8_t> png_bytes;
    if (!i3slib::utl::encode_png(
          reinterpret_cast<uint8_t*>(image->data()), width, height, false, png_bytes))
      return false;
    texture_buffer.meta.alpha_status = i3slib::i3s::Texture_meta::Alpha_status::Opaque;
    texture_buffer.meta.wrap_mode = i3slib::i3s::Texture_meta::Wrap_mode::None;
    texture_buffer.meta.is_atlas = false;
    texture_buffer.meta.mip0_width = width;
    texture_buffer.meta.mip0_height = height;
    texture_buffer.meta.mip_count = 1; //no mip
    texture_buffer.meta.format = i3slib::i3s::Image_format::Png;

    texture_buffer.data = i3slib::utl::Buffer::create_deep_copy<char>(
      reinterpret_cast<const char*>(png_bytes.data()),
      static_cast<int>(png_bytes.size()));
    return true;
  }

  // Extracts the fragment of the elevation grid starting at _start_ and spanning _size_ cells 
  // (i.e, size + 1 nodes) and extracts the fragment of color_data starting at texture_start
  // with texture_size * texture_size pixels.

  bool build_mesh(
    const i3slib::i3s::Layer_writer& writer,
    osg::Geometry* geometry,
    i3slib::i3s::Mesh_data& mesh,
    CS_transformation* transformation = nullptr)
  {
    std::vector<Vec3d> verts;
    std::vector<Vec2f> uvs;

    // 获取顶点坐标
    osg::Vec3Array* osg_verts = dynamic_cast<osg::Vec3Array*>(geometry->getVertexArray());
    //获取UV
    osg::Vec2Array* osg_vertsUV = dynamic_cast<osg::Vec2Array*>(geometry->getTexCoordArray(0));

    // 获取图元
    // osg::Geometry::PrimitiveSetList& psets = geometry->getPrimitiveSetList();
    int numP = geometry->getNumPrimitiveSets();
    for (unsigned int ipr = 0; ipr < numP; ipr++)
    {
        osg::PrimitiveSet* prset = geometry->getPrimitiveSet(ipr);
        bool find = false;
        if(prset->getMode() == GL_TRIANGLES)
        {
          find = true;
        }
        if(find)
        {
          unsigned int ncnt = prset->getNumIndices();
          for (unsigned int ic = 0; ic * 3 < ncnt; ic++)
          {
              unsigned int iIndex0 = prset->index(ic * 3);
              unsigned int iIndex1 = prset->index(ic * 3 + 1);
              unsigned int iIndex2 = prset->index(ic * 3 + 2);
              osg::Vec3 osg_v0 = osg_verts->at(iIndex0);
              osg::Vec3 osg_v1 = osg_verts->at(iIndex1);
              osg::Vec3 osg_v2 = osg_verts->at(iIndex2);
              const Vec3d v0(osg_v0.x(), osg_v0.y(), osg_v0.z());
              const Vec3d v1(osg_v1.x(), osg_v1.y(), osg_v1.z());
              const Vec3d v2(osg_v2.x(), osg_v2.y(), osg_v2.z());
              verts.push_back(v0);
              verts.push_back(v1);
              verts.push_back(v2);
              if(osg_vertsUV){
                osg::Vec2 osg_u0 = osg_vertsUV->at(iIndex0);
                osg::Vec2 osg_u1 = osg_vertsUV->at(iIndex1);
                osg::Vec2 osg_u2 = osg_vertsUV->at(iIndex2);
                const Vec2f u0(osg_u0.x(), osg_u0.y());
                const Vec2f u1(osg_u1.x(), osg_u1.y());
                const Vec2f u2(osg_u2.x(), osg_u2.y());
                uvs.push_back(clamp_uv(u0));
                uvs.push_back(clamp_uv(u1));
                uvs.push_back(clamp_uv(u2));
              }
          }
        }
        else{
            printf(" prset->getMode() not GL_TRIANGLES \n");
        }
    }

    if (transformation)
      transformation->transform(verts.data(), verts.size());

    i3slib::i3s::Simple_raw_mesh raw_mesh;
    raw_mesh.vertex_count = static_cast<int>(verts.size());
    raw_mesh.abs_xyz = verts.data();
    raw_mesh.uv = uvs.data();
    osg::Texture2D* tex2D = dynamic_cast<osg::Texture2D*>(geometry->getStateSet()->getTextureAttribute(0, osg::StateAttribute::TEXTURE));
    if(tex2D){
      if (!extract_texture_fragment(tex2D, raw_mesh.img)) return false;
    }
    return writer.create_mesh_from_raw(raw_mesh, mesh) == IDS_I3S_OK;
  }

  void tile_process(i3slib::i3s::Layer_writer& writer, std::string& file_name, int depth, i3slib::i3s::Node_id& node_id, CS_transformation* transformation = nullptr) {
    std::vector<std::string> fileNames = { file_name };
    std::string parent_path = get_parent(file_name);
    osg::ref_ptr<osg::Node> root = osgDB::readNodeFiles(fileNames);
    if (!root.valid()) {
      std::cerr << "read node files " << file_name << " fail! \n";
      return;
    }
    InfoVisitor infoVisitor(parent_path);
    root->accept(infoVisitor);
    if (infoVisitor.geometry_array.empty()) return;
    std::vector<i3slib::i3s::Node_id> node_parent_ids;
    int num = infoVisitor.geometry_array.size();
    for (int d = 0; d < num; d++)
    {
      std::vector<i3slib::i3s::Node_id> child_ids;
      auto geometry = infoVisitor.geometry_array[d];
      // 处理node
      i3slib::i3s::Simple_node_data node_data;
      node_data.node_depth = depth + 1;
      node_data.children = std::move(child_ids);
      node_data.lod_threshold = screen_size_to_area(500);
      const auto status = build_mesh(writer, geometry, node_data.mesh, transformation);
      if (!status)
        return;
      std::cout << "create_node:node_id" << node_id << ",node_depth:"<< node_data.node_depth << std::endl;
      if (writer.create_node(node_data, node_id) != IDS_I3S_OK)
        return;
      node_parent_ids.push_back(node_id++);
    }  
    // Add a root node on top of everything.
    i3slib::i3s::Simple_node_data node_parent_data;
    node_parent_data.node_depth = depth;
    node_parent_data.children = std::move(node_parent_ids);
    std::cout << "create_parent_node:node_id" << node_id << ",node_depth:"<< node_parent_data.node_depth<< ",size:"<< node_parent_data.children.size()<< std::endl;
    if (writer.create_node(node_parent_data, node_id) != IDS_I3S_OK)
      return; 
  }
  void get_all_tree(i3slib::i3s::Layer_writer& writer, std::string& file_name, int depth, i3slib::i3s::Node_id& node_id, CS_transformation* transformation = nullptr) {
    std::vector<std::string> fileNames = { file_name };
    InfoVisitor infoVisitor(get_parent(file_name));
    {   // add block to release Node
        osg::ref_ptr<osg::Node> root = osgDB::readNodeFiles(fileNames);
        if (!root) {
            std::cerr << "read node files " << file_name << " fail! \n";
            return ;
        }
        root->accept(infoVisitor);    
    }
    if(infoVisitor.sub_node_names.size() > 0){
      std::vector<i3slib::i3s::Node_id> node_ids;
      for (auto& i : infoVisitor.sub_node_names) {
        get_all_tree(writer,i,depth + 1,node_id,transformation);
        if (!file_name.empty()) {
            node_ids.push_back(node_id++);
        }
      }
      // Add a root node on top of everything.
      i3slib::i3s::Simple_node_data node_parent_data;
      node_parent_data.node_depth = depth;
      node_parent_data.children = std::move(node_ids);
      if (writer.create_node(node_parent_data, node_id) != IDS_I3S_OK)
        return; 
      std::cout << "create_parent_node:node_id" << node_id << ",node_depth:"<< depth << std::endl;
    }else{
      if (infoVisitor.geometry_array.empty()) return;
      tile_process(writer,file_name,depth+1,node_id,transformation);
    }
  }

  void process(std::string source_dir_, std::string target_dir_, CS_transformation* transformation = nullptr){
    for (const auto &data_dir: stdfs::directory_iterator(source_dir_,
            ~stdfs::directory_options::follow_directory_symlink
            | stdfs::directory_options::skip_permission_denied)) {
        if(!stdfs::is_directory(data_dir.path()))
            continue; 
        std::string tile_name = data_dir.path().stem().string();
        std::string tile_path = (data_dir.path() / (tile_name + ".osgb")).string();
        if(!stdfs::exists(tile_path)){
            process(data_dir.path().string(),(stdfs::path(target_dir_) / stdfs::path(tile_name)).string(),transformation);
            continue;
        }
        std::cout << "find data " << data_dir<< ", start process tile" << std::endl;
        std::string tile_dir = data_dir.path().string();
        std::string out_dir = (stdfs::path(target_dir_) / stdfs::path(tile_name)).string();

        if (!stdfs::exists(target_dir_)) {
            stdfs::create_directories(target_dir_);
        }
        auto writer = create_writer(stdfs::path(out_dir + ".slpk"));
        if (!writer)
          return;
        i3slib::i3s::Node_id node_id = 0;
        get_all_tree(*writer,tile_path,0,node_id,transformation);
        if (writer->save() != IDS_I3S_OK)
          return;
        std::cout << "data " << data_dir<< ", export slpk success" << std::endl;
    }
  }
}

int main(int argc, char* argv[])
{
  if (argc < 3)
  {
    std::cout << "Usage:" << std::endl
      << "lod2slpk <source_osgb_file> <output_slpk_file>" << std::endl;
    return 1;
  }

  ENU_to_WGS_transformation transformation({ 119.0133760800, 25.4370228400 });

  std::string source_dir_ = std::string( argv[1]) ;
  std::string target_dir_ = std::string( argv[2]) ;
  if (source_dir_.back() == '/') {
      source_dir_.pop_back(); 
  }
  std::string model_file = stdfs::path(source_dir_).stem().string();
  std::string out_dir = (stdfs::path(target_dir_) / stdfs::path(model_file)).string();
  if (!stdfs::exists(out_dir)) {
      stdfs::create_directories(out_dir);
  }
  process(source_dir_,out_dir,&transformation);

  std::cout << "end process" << std::endl;
}
